#!/usr/bin/env python
# -*- conding: utf-8 -*-

"""
@Time     : 2024/8/28 6:57
@Author   : liujingmao
@File     : 01.CacheBackEmbedding使用示例.py
"""

import dotenv
import numpy as np
from langchain.embeddings import CacheBackedEmbeddings
from langchain.storage import LocalFileStore
from langchain_openai import OpenAIEmbeddings
from numpy.linalg import norm

dotenv.load_dotenv()


# 计算consine相似度
def consine_similarity(vector1: list, vector2: list) -> float:
    """计算传入两个向量的余弦相似度"""
    # 1. 计算两个向量的点积
    dot_product = np.dot(vector1, vector2)

    # 2. 计算向量的长度
    vec1_norm = norm(vector1)
    vec2_norm = norm(vector2)

    # 3. 计算余弦相似度
    return dot_product / (vec1_norm * vec2_norm)


# 1. 创建文本嵌入模型
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

embeddings_with_cache = CacheBackedEmbeddings.from_bytes_store(
    embeddings,
    LocalFileStore("./cache/"),
    namespace=embeddings.model,
    query_embedding_cache=True,
)

# 2. 嵌入文本
query_vector = embeddings_with_cache.embed_query("我是茂茂，我喜欢篮球")

# print(query_vector)
print(len(query_vector))
print(type(query_vector))

# 3. 嵌入文档列表/字符串列表
document_vector = embeddings_with_cache.embed_documents(
    [
        "我叫慕小课，我喜欢打篮球",
        "这个喜欢打篮球的人叫慕小课",
        "慕小课喜欢打篮球"
    ]
)
print(len(document_vector))

# 4. 计算余弦相似度
print("向量1和向量2的相似度", consine_similarity(document_vector[0], document_vector[1]))
print("向量1和向量3的相似度", consine_similarity(document_vector[0], document_vector[2]))
